# -*- coding:utf-8 -*-
import numpy as np


class HardTemplate(object):
    """
    硬模板，人工定义句子和[MASK]之间的位置关系。
    """

    def __init__(self, prompt: str):
        """
        Args:
            prompt (str): prompt格式定义字符串, e.g. -> "这是一条{MASK}评论：{textA}。"
        """
        self.prompt = prompt
        self.inputs_list = []  # 根据文字prompt拆分为各part的列表
        self.custom_tokens = set(['MASK'])  # 从prompt中解析出的自定义token集合
        self.prompt_analysis()  # 解析prompt模板

    def prompt_analysis(self):
        """
        将prompt文字模板拆解为可映射的数据结构。

        Examples:
            prompt -> "这是一条{MASK}评论：{textA}。"
            inputs_list -> ['这', '是', '一', '条', 'MASK', '评', '论', '：', 'textA', '。']
            custom_tokens -> {'textA', 'MASK'}
        """
        idx = 0
        while idx < len(self.prompt):
            str_part = ''
            if self.prompt[idx] not in ['{', '}']:
                self.inputs_list.append(self.prompt[idx])
            if self.prompt[idx] == '{':  # 进入自定义字段
                idx += 1
                while self.prompt[idx] != '}':
                    str_part += self.prompt[idx]  # 拼接该自定义字段的值
                    idx += 1
            elif self.prompt[idx] == '}':
                raise ValueError("Unmatched bracket '}', check your prompt.")
            if str_part:
                self.inputs_list.append(str_part)
                # 将所有自定义字段存储，后续会检测输入信息是否完整
                self.custom_tokens.add(str_part)
            idx += 1

    def __call__(self,
                 inputs_dict: dict,
                 tokenizer,
                 mask_length,
                 max_seq_len=512):
        """
        输入一个样本，转换为符合模板的格式。

        Args:
            inputs_dict (dict): prompt中的参数字典, e.g. -> {
                                                            "textA": "这个手机也太卡了", 
                                                            "MASK": "[MASK]"
                                                        }
            tokenizer: 用于encoding文本
            mask_length (int): MASK token 的长度
            max_seq_len : 最大句子长度

        Returns:
            dict -> {
                'text': '[CLS]这是一条[MASK]评论：这个手机也太卡了。[SEP]',
                'input_ids': [1, 47, 10, 7, 304, 3, 480, 279, 74, 47, 27, 247, 98, 105, 512, 777, 15, 12043, 2],
                'token_type_ids': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], 
                'attention_mask': [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
                'mask_position': [0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
            }
        """
        # 定义输出格式
        outputs = {
            'text': '',
            'input_ids': [],
            'token_type_ids': [],
            'attention_mask': [],
            'mask_position': []
        }

        text = ''
        for value in self.inputs_list:
            if value in self.custom_tokens:
                if value == 'MASK':
                    text += inputs_dict[value] * mask_length
                else:
                    text += inputs_dict[value]
            else:
                text += value
        # print(f'text-->{text}')
        encoded = tokenizer(text=text,
                            truncation=True,  # 启用截断（超过 max_length 的文本被截断）
                            max_length=max_seq_len,  # 最大序列长度（含特殊标记）
                            padding="max_length")
        # 填充策略（padding）
        # "max_length"：强制填充到 max_length
        # "longest"：按批次中最长序列填充（需配合 batched=True）
        # False：不填充（可能导致批次内长度不一致）

        # print(f'encoded--->{encoded}')
        outputs['input_ids'] = encoded['input_ids']
        # token_type_ids 的适用性：
        # 仅当模型需要区分双句子输入（如BERT的NSP任务）时才会生成。
        outputs['token_type_ids'] = encoded['token_type_ids']

        # 标记哪些位置是有效 token（1）或填充位（0）
        outputs['attention_mask'] = encoded['attention_mask']

        outputs['text'] = ''.join(tokenizer.convert_ids_to_tokens(encoded['input_ids']))
        mask_token_id = tokenizer.convert_tokens_to_ids(['[MASK]'])[0]
        mask_position = np.where(np.array(outputs['input_ids']) == mask_token_id)[0].tolist()
        outputs['mask_position'] = mask_position
        return outputs
